local Class = require('lua.lang.Class')
local BufferedInputStream = require('lua.io.BufferedInputStream')
local RuntimeException = require('lua.lang.RuntimeException')
local TokenKind = require('lua.parser.TokenKind')
local LexToken = require('lua.parser.LexToken')

local B_SPACE = (' '):byte()
local B_CR = ('\r'):byte()
local B_LF = ('\n'):byte()
local B_TAB = ('\t'):byte()
local B_HASH = ('#'):byte()
local B_EXCLAMATION = ('!'):byte()
local B_LBRACKET = ('['):byte()
local B_RBRACKET = (']'):byte()
local B_EQUALS = ('='):byte()
local B_MINUS = ('-'):byte()
local B_PLUS = ('+'):byte()
local B_a = ('a'):byte()
local B_A = ('A'):byte()
local B_z = ('z'):byte()
local B_Z = ('Z'):byte()
local B_UNDERSCORE = ('_'):byte()
local B_ZERO = ('0'):byte()
local B_NINE = ('9'):byte()
local B_DOT = ('.'):byte()
local B_x = ('x'):byte()
local B_X = ('X'):byte()
local B_p = ('p'):byte()
local B_P = ('P'):byte()
local B_e = ('e'):byte()
local B_E = ('E'):byte()
local B_f = ('f'):byte()
local B_F = ('F'):byte()
local B_BACKSLASH = ('\\'):byte()
local B_LESS_THAN = ('<'):byte()
local B_GREATER_THAN = ('>'):byte()
local B_TILDE = ('~'):byte()
local B_SLASH = ('/'):byte()
local B_AMPERSAND = ('&'):byte()
local B_COLON = (':'):byte()

--- @class Lexer : Class
--- @field is InputStream
--- @field line number
--- @field column number
--- @field buff number[]
local Lexer = Class {}

--- @type table<string, boolean>
Lexer.MAP_KEYWORDS = {
	['and'] = true,
	['break'] = true,
	['do'] = true,
	['else'] = true,
	['elseif'] = true,
	['end'] = true,
	['false'] = true,
	['for'] = true,
	['function'] = true,
	['if'] = true,
	['in'] = true,
	['local'] = true,
	['nil'] = true,
	['not'] = true,
	['or'] = true,
	['repeat'] = true,
	['return'] = true,
	['then'] = true,
	['true'] = true,
	['until'] = true,
	['while'] = true
}

--- @type table<number, boolean>
Lexer.MAP_SYMBOLS = {
	[B_PLUS] = true,
	[B_MINUS] = true,
	[('*'):byte()] = true,
	[B_SLASH] = true,
	[('^'):byte()] = true,
	[('%'):byte()] = true,
	[(','):byte()] = true,
	[('{'):byte()] = true,
	[('}'):byte()] = true,
	[B_LBRACKET] = true,
	[B_RBRACKET] = true,
	[('('):byte()] = true,
	[(')'):byte()] = true,
	[(';'):byte()] = true,
	[B_HASH] = true,
	[('|'):byte()] = true,
	[B_AMPERSAND] = true,
	[B_TILDE] = true
}

--- @param is InputStream
--- @return Lexer
function Lexer:new(is)
	--- @type Lexer
	local o = {}
	if not is:markSupported() then
		is = BufferedInputStream:new(is)
	end
	o.is = is
	o.line = 1
	o.column = 1
	o.buff = {}
	return setmetatable(o, Lexer)
end

--- @return number
function Lexer:get()
	local c = self.is:read()
	if c == B_LF then
		self.column = 1
		self.line = self.line + 1
	else
		self.column = self.column + 1
	end
	return c
end

--- @param n number?
--- @return number
function Lexer:peek(n)
	n = n or 0
	local is = self.is
	is:mark(n)
	local c = self:get()
	for _ = 1, n do
		c = self:get()
	end
	is:reset()
	return c
end

--- @param chars number[]
--- @return number
function Lexer:consume(chars)
	local c = self:peek()
	for i = 1, #chars do
		if chars[i] == c then
			return self:getAndSave()
		end
	end
	return -1
end

--- @param errMsg string
--- @return void
function Lexer:generateError(errMsg)
	RuntimeException:new(('>> :%s:%s: %s'):format(self.line, self.column, errMsg)):throw()
end

--- @return number
function Lexer:getAndSave()
	local buff = self.buff
	local c = self:get()
	buff[#buff + 1] = c
	return c
end

--- @return string
function Lexer:reset()
	local s = string.char(table.unpack(self.buff))
	self.buff = {}
	return s
end

--- @return boolean
function Lexer:tryGetLongString()
	if self:peek() ~= B_LBRACKET then
		return false
	end
	local equalsCount = 0
	local depth = 1
	while self:peek(equalsCount + 1) == B_EQUALS do
		equalsCount = equalsCount + 1
	end
	if self:peek(equalsCount + 1) ~= B_LBRACKET then
		return false
	end
	for _ = 1, equalsCount + 2 do
		self:getAndSave()
	end
	while true do
		if self:peek() == -1 then
			self:generateError('Expected `]' .. ('='):rep(equalsCount) .. ']` near <eof>.')
		end
		local foundEnd = true
		if self:peek() == B_RBRACKET then
			for i = 1, equalsCount do
				if self:peek(i) ~= B_EQUALS then
					foundEnd = false
					break
				end
			end
			if self:peek(equalsCount + 1) ~= B_RBRACKET then
				foundEnd = false
			end
		else
			if self:peek() == B_LBRACKET then
				local embedded = true
				for i = 1, equalsCount do
					if self:peek(i) ~= B_EQUALS then
						embedded = false
						break
					end
				end
				if self:peek(equalsCount + 1) == B_LBRACKET and embedded then
					depth = depth + 1
					for _ = 1, equalsCount + 2 do
						self:getAndSave()
					end
				end
			end
			foundEnd = false
		end
		if foundEnd then
			depth = depth - 1
			if depth == 0 then
				break
			else
				for _ = 1, equalsCount + 2 do
					self:getAndSave()
				end
			end
		else
			self:getAndSave()
		end
	end
	for _ = 1, equalsCount + 2 do
		self:getAndSave()
	end
	return true
end

--- @param c number
--- @return boolean
function Lexer:isAlpha(c)
	return (c >= B_A and c <= B_Z) or (c >= B_a and c <= B_z)
end

--- @param c number
--- @return boolean
function Lexer:isDigit(c)
	return c >= B_ZERO and c <= B_NINE
end

--- @param c number
--- @return boolean
function Lexer:isXDigit(c)
	return (c >= B_ZERO and c <= B_NINE) or (c >= B_A and c <= B_F) or (c >= B_a and c <= B_f)
end

--- @return LexToken?
function Lexer:getToken()
	local c = self:peek()
	local line = self.line
	local column = self.column
	if c == B_HASH and self:peek(1) == B_EXCLAMATION and self.line == 1 then
		self:getAndSave()
		self:getAndSave()
		while self:peek() ~= B_LF and self:peek() ~= -1 do
			self:getAndSave()
		end
		return LexToken:new(TokenKind.Comment, line, column, self:reset())
	elseif c == B_TAB or c == B_SPACE or c == B_LF or c == B_CR then
		self:getAndSave()
		while self:peek() == B_TAB or self:peek() == B_SPACE or
				self:peek() == B_LF or self:peek() == B_CR do
			self:getAndSave()
		end
		return LexToken:new(TokenKind.WhiteSpace, line, column, self:reset())
	elseif c == B_MINUS and self:peek(1) == B_MINUS then
		self:getAndSave()
		self:getAndSave()
		if not self:tryGetLongString() then
			while self:peek() ~= B_LF and self:peek() ~= -1 do
				self:getAndSave()
			end
		end
		return LexToken:new(TokenKind.Comment, line, column, self:reset())
	elseif c == -1 then
		return LexToken:new(TokenKind.Eof, line, column, '')
	elseif self:isAlpha(c) or c == B_UNDERSCORE then
		repeat
			self:getAndSave()
			c = self:peek()
		until not (self:isAlpha(c) or c == B_UNDERSCORE or self:isDigit(c))
		local s = self:reset()
		return LexToken:new(self.MAP_KEYWORDS[s] and TokenKind.KeyWord or TokenKind.Identifier, line, column, s)
	elseif self:isDigit(c) or (c == B_DOT and self:isDigit(self:peek(1))) then
		local d = self:peek(1)
		if c == B_ZERO and (d == B_x or d == B_X) then
			self:getAndSave()
			self:getAndSave()
			while self:isXDigit(self:peek()) do
				self:getAndSave()
			end
			if self:consume({ B_P, B_p }) ~= -1 then
				self:consume({ B_PLUS, B_MINUS })
				while self:isDigit(self:peek()) do
					self:getAndSave()
				end
			end
		else
			while self:isDigit(self:peek()) do
				self:getAndSave()
			end
			if self:consume({ B_DOT }) ~= -1 then
				while self:isDigit(self:peek()) do
					self:getAndSave()
				end
			end
			if self:consume({ B_E, B_e }) ~= -1  then
				self:consume({ B_PLUS, B_MINUS })
				while self:isDigit(self:peek()) do
					self:getAndSave()
				end
			end
		end
		return LexToken:new(TokenKind.Number, line, column, self:reset())
	elseif c == 39 or c == 34 then
		local delim = self:getAndSave()
		while true do
			c = self:getAndSave()
			if c == B_BACKSLASH then
				self:getAndSave()
			elseif c == delim then
				break
			elseif c == -1 then
				self:generateError('Unfinished string `' .. string.char(delim) .. '` near <eof>.')
			end
		end
		return LexToken:new(TokenKind.String, line, column, self:reset())
	elseif c == B_LBRACKET then
		if self:tryGetLongString() then
			return LexToken:new(TokenKind.String, line, column, self:reset())
		else
			self:getAndSave()
			return LexToken:new(TokenKind.Symbol, line, column, self:reset())
		end
	elseif self:consume({ B_LESS_THAN, B_EQUALS, B_GREATER_THAN }) ~= -1 then
		if self:consume({ B_EQUALS }) == -1 then
			self:consume({ c })
		end
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self:consume({ B_TILDE }) ~= -1 then
		self:consume({ B_EQUALS })
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self:consume({ B_SLASH }) ~= -1 then
		self:consume({ B_SLASH })
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self:consume({ B_AMPERSAND }) ~= -1 then
		self:consume({ B_AMPERSAND, B_EQUALS })
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self:consume({ B_DOT }) ~= -1 then
		self:consume({ B_DOT })
		self:consume({ B_DOT })
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self:consume({ B_COLON }) ~= -1 then
		self:consume({ B_COLON })
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	elseif self.MAP_SYMBOLS[c] then
		self:getAndSave()
		return LexToken:new(TokenKind.Symbol, line, column, self:reset())
	else
		self:generateError('Unexpected Symbol `' .. string.char(c) .. '` in source.')
	end
	return nil
end

return Lexer